{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Small Language Model Offline Serving with vLLM\n",
        "This notebook is a companion of chapter 10 of the \"Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2024.  \n",
        "The code in this notebook is about serving some example small language models, such as [GPT-2 small](https://huggingface.co/openai-community/gpt2) and [Microsoft's Phi-3 mini 4k Instruct](microsoft/Phi-3-mini-4k-instruct), through the [vLLM](https://github.com/vllm-project/vllm) API. The same code applies to any other Open Source LLM. Hardware acceleration is needed.  \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "BEPpASlVBHg8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the vLLM package. Installing through `pip` requires the CUDA drivers and an NVIDIA GPU, as there is no pre-built wheel for CPU only available for Linux systems. If you want to run the code below for GPT-2 small without hardware acceleration you have to build vLLM from source instead. Please see the official vLLM documentation for steps on how to do it.  \n",
        "Restart the VM once the vLLM installation completes."
      ],
      "metadata": {
        "id": "_AgHqzpHCssF"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SOvu_X8wVr3z"
      },
      "outputs": [],
      "source": [
        "!pip install vllm"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Text Generation"
      ],
      "metadata": {
        "id": "Jtj6NZgYC3-B"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the required packages and classes."
      ],
      "metadata": {
        "id": "kdh7i-r0CvV2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from vllm import LLM"
      ],
      "metadata": {
        "id": "vJYOjm2UVy2M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Specify the model to use (GPT-2 small in this case) and where to offload its weights (CPU or GPU) and then create an instance of the `vLLM` engine class for it. This class, after downloading the configuration file for the given model from the HF's Hub, analyzes the configuration and prepare the model setup accordingly, before downloading checkpoitns and tokenizer. Once the download is completed, it laodt the checkpoints in the destinazion device, profile the memory, initializes the KV cache and warms up the model."
      ],
      "metadata": {
        "id": "pVicz6eHCzGF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_id = \"openai-community/gpt2\"\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "vllm_model = LLM(model_id, device=device, enforce_eager=True)"
      ],
      "metadata": {
        "id": "cf_kfnPRVz8S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Provide a prompt."
      ],
      "metadata": {
        "id": "5BguPrYCEvy3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"Once upon a time in a land far away\""
      ],
      "metadata": {
        "id": "Hr-HsXuxV6Y7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implement a custom function to generate text with the given model in vLLM."
      ],
      "metadata": {
        "id": "v8gSjFmSE04I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_text(prompt, max_length=50):\n",
        "    input_ids = vllm_model.tokenizer.encode(prompt,\n",
        "                                            return_tensors=\"pt\").to(device)\n",
        "\n",
        "    generated_ids = vllm_model.generate(\n",
        "        input_ids=input_ids,\n",
        "        max_length=max_length,\n",
        "        do_sample=True,\n",
        "        num_return_sequences=1\n",
        "    )\n",
        "\n",
        "    generated_text = vllm_model.tokenizer.decode(generated_ids[0],\n",
        "                                                 skip_special_tokens=True)\n",
        "\n",
        "    return generated_text"
      ],
      "metadata": {
        "id": "XFqXBP9zV9P4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Call the `generate_text` function and display the generated text."
      ],
      "metadata": {
        "id": "s9tlijmaFF93"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "output = vllm_model.generate(prompt)\n",
        "for item in output:\n",
        "    prompt = item.prompt\n",
        "    generated_text = item.outputs[0].text\n",
        "    print(generated_text)"
      ],
      "metadata": {
        "id": "d27kPl-gWGId"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Chat"
      ],
      "metadata": {
        "id": "aE3C0X1mI-XH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del vllm_model\n",
        "\n",
        "import gc\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "2x9x-055VVeB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Specify the model to use (Microsoft Phi-3 mini 4k Instruct in this case) and where to offload its weights (CPU or GPU) and then create an instance of the `vLLM` engine class for it. Also using the `SamplingParams` class to set the temperature for the model."
      ],
      "metadata": {
        "id": "qK4EHQnUIE9t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from vllm import LLM, SamplingParams\n",
        "\n",
        "model_id = \"microsoft/Phi-3-mini-4k-instruct\"\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "llm = LLM(model=model_id, device=device, dtype=torch.half)\n",
        "\n",
        "sampling_params = SamplingParams(temperature=0.5)"
      ],
      "metadata": {
        "id": "gpHb1WlOdF0m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a function to print the generated outputs."
      ],
      "metadata": {
        "id": "2bpYFrCEOxEh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def print_outputs(outputs):\n",
        "  print(\"=\" * 80)\n",
        "  for output in outputs:\n",
        "      prompt = output.prompt\n",
        "      generated_text = output.outputs[0].text\n",
        "      print(f\"Prompt: {prompt!r}, Generated text: {generated_text!r}\")\n",
        "  print(\"-\" * 80)"
      ],
      "metadata": {
        "id": "8mHdA5evbfCz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Setup the input in a format to chat with the selected model, then run the `chat` command and display the output."
      ],
      "metadata": {
        "id": "HM7fvvMWO1qd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "conversation = [\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": \"You are a helpful assistant\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"Hi\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"assistant\",\n",
        "        \"content\": \"Hi! How can I assist you today?\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"Write an essay about the Monte Carlo Tree Search algorithm.\",\n",
        "    },\n",
        "]\n",
        "outputs = llm.chat(conversation,\n",
        "                   sampling_params=sampling_params,\n",
        "                   use_tqdm=False)\n",
        "print_outputs(outputs)"
      ],
      "metadata": {
        "id": "iMyFmU0OIal8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The same way, batch inference can be run with the vLLM's chat API."
      ],
      "metadata": {
        "id": "xGabTGefPGTk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "conversation = [\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": \"You are a helpful assistant\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"Hi\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"assistant\",\n",
        "        \"content\": \"Hi! How can I assist you today?\"\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"Write an essay about the Monte Carlo Tree Search algorithm.\",\n",
        "    },\n",
        "]\n",
        "conversations = [conversation for _ in range(10)]\n",
        "\n",
        "outputs = llm.chat(messages=conversations,\n",
        "                   sampling_params=sampling_params,\n",
        "                   use_tqdm=True)\n",
        "print_outputs(outputs)"
      ],
      "metadata": {
        "id": "cpdy7PC7bm3E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "It is also possible to specify a chat template to the vLLM's chat API. Download one for the Phi-3 mini model."
      ],
      "metadata": {
        "id": "coBJDC_tPYGo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/chujiezheng/chat_templates/main/chat_templates/phi-3.jinja"
      ],
      "metadata": {
        "id": "vBNgVk_DgGUx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then supply it to the inference engine class for running chats."
      ],
      "metadata": {
        "id": "PhXwCsK2PiVU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with open('phi-3.jinja', \"r\") as f:\n",
        "  chat_template = f.read()\n",
        "\n",
        "outputs = llm.chat(\n",
        "  conversations,\n",
        "  sampling_params=sampling_params,\n",
        "  use_tqdm=False,\n",
        "  chat_template=chat_template,\n",
        ")"
      ],
      "metadata": {
        "id": "7am14tC8f6bF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "outputs"
      ],
      "metadata": {
        "id": "IDUsXnvAgs9l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# GGUF"
      ],
      "metadata": {
        "id": "L430i0c44BnP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "del llm\n",
        "\n",
        "import gc\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "PniFmx9zX3nb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "vLLM supports also models in [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) format. The code example in the cells below show evidence of this using a [GGUF version](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct-gguf) of the Phi-3 mini 4k Instruct model already available from Microsoft in the HF's Hub. Download it and the tokenizer associated to the vanilla model too."
      ],
      "metadata": {
        "id": "r7VLcj5sPxns"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "repo_id = \"microsoft/Phi-3-mini-4k-instruct-gguf\"\n",
        "filename = \"Phi-3-mini-4k-instruct-q4.gguf\"\n",
        "tokenizer = \"microsoft/Phi-3-mini-4k-instruct\"\n",
        "model = hf_hub_download(repo_id, filename=filename)"
      ],
      "metadata": {
        "id": "nI3vqk5J4Cr_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before starting the chat, uninstall the PyTorch-native training-to-serving model optimization package (lately available by default in the Colab VM), as it interferes with the vLLM GGUF inference."
      ],
      "metadata": {
        "id": "JSIEgooqgcrz"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d0457d58"
      },
      "source": [
        "!pip uninstall -y torchao"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Once the download is completed, create an instance of the vLLM engine starting from the model in the local file system. Also setup some prompt to chat with it and configure some sampling parameters (`temperature` and `max_tokens`)."
      ],
      "metadata": {
        "id": "ltKIyI28RH-x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from vllm import LLM, SamplingParams\n",
        "\n",
        "prompts = [\n",
        "    \"How to explain Internet for a medieval knight?\",\n",
        "    \"What's the future of AI?\",\n",
        "]\n",
        "prompts = [[{\"role\": \"user\", \"content\": prompt}] for prompt in prompts]\n",
        "sampling_params = SamplingParams(temperature=0, max_tokens=128)\n",
        "\n",
        "llm = LLM(model=model, tokenizer=tokenizer)"
      ],
      "metadata": {
        "id": "8iJyTnC0i0Bk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Start the chat with the GGUF model and display the outputs."
      ],
      "metadata": {
        "id": "1gXlZgDFRxlA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "outputs = llm.chat(prompts, sampling_params)\n",
        "print_outputs(outputs)"
      ],
      "metadata": {
        "id": "KONcVGv-5ByB"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}